1882D - Tree XOR - CodeForces Solution

#include<bits/stdc++.h>

using namespace std;
#define ll long long
#define endl "\n"
#define momen ios_base::sync_with_stdio(0);cin.tie(0);cout.tie(0);
#define fp(n) for(int i=0;i<n;i++)
#define fp1(n) for(int i=1;i<=n;i++)
#define all(v) v.begin() , v.end()
#define matrix vector<vector<ll>>
//#define f first
//#define s second
#define bk back()
#define pb push_back()
#define sz(x) int((x).size())

const int N = 2e5+5  , mod = 998244353 ;
int arr[N] ;
int n;

vector<int>adj[N] ;

ll ans[N] , subT[N] , ret[N] ;


ll dfs(int i , int p){
    for(auto j:adj[i]){
        if (p == j)
            continue;
        ret[i] += dfs(j ,i  ) ;
    }
    for(auto j:adj[i]){
        if (p == j)
            continue;
    }


    return ret[i] + subT[i] *(arr[i]^arr[p]) ;
}

void dfs1(int i ,int p ,ll up){

    ans[i] = ret[i] + up ;
    for(auto j:adj[i]){
        if (p == j)
            continue;
        up+= ret[j] +  subT[j] *(arr[j]^arr[i])  ;
    }
    for(auto j :adj[i]){
        if (j == p)
            continue;
        ll x = (n - subT[j]) * (arr[i] ^ arr[j]) ;
        dfs1(j , i , up - ret[j] -  subT[j] *(arr[j]^arr[i])  + x ) ;
    }
}

int dfs2(int i , int p){
    int x = 1;
    for(auto j:adj[i]){
        if(j == p)
            continue;
        x += dfs2(j ,i  ) ;
    }
    subT[i] = x ;
    return x ;
}

void solve() {

    cin>>n;
    for (int i = 1; i <= n; ++i) {
        cin>>arr[i] ;
        ans[i]=  0 ;
        ret[i] = 0;
        subT[i] = 0 ;
        adj[i].clear() ;
    }
    for (int i = 0; i < n-1; ++i) {
        int a,b ;
        cin>>a>>b;
        adj[a].push_back(b) ;
        adj[b].push_back(a) ;
    }
    dfs2(1,0);
    dfs(1,0) ;
    dfs1(1,0,0) ;
    for (int i = 1; i <= n ; ++i) {
        cout<<ans[i]<<" ";
    }
    cout<<endl;
}


signed main() {

    momen

    int t = 1;
#ifndef ONLINE_JUDGE
    freopen("input.txt", "r", stdin);
    freopen("output.txt", "w", stdout);
#endif
    cin >> t;
    for (int i = 1; i <= t; ++i) {
//        cout<<"Case #"<<i<<": ";
        solve();
    }

    return 0;
}